import os
import xarray as xr
import numpy as np
import rasterio
from rasterio.transform import from_origin

def calculate_daily_mean_precipitation(input_folder, output_folder, block_size=64):
    # 获取文件夹中的所有.nc文件
    nc_files = [f for f in os.listdir(input_folder) if f.endswith('.nc')]
    
    for nc_file in nc_files:
        # 读取每一个NetCDF文件
        nc_path = os.path.join(input_folder, nc_file)
        ds = xr.open_dataset(nc_path, chunks={'time': 1})  # 按时间维度逐天分块读取数据
        
        if 'SMCI' not in ds.variables:
            print(f"文件 {nc_file} 中未找到变量 'SMCI'")
            continue

        time_len= ds.dims['time']
        print(f"\n文件 {nc_file} 中共有 {time_len}个时间步")

        # 计算日均降水量，跳过缺失值并按块处理，避免占用过多内存
        daily_mean_precip = ds['SMCI'].mean(dim='time', skipna=True)

        print(f"{nc_file}已完成日均降水量计算\n")
        # daily_mean_precip=None
        # for t,step in enumerate(ds['SMCI'],start=1):
        #     if daily_mean_precip is None:
        #         daily_mean_precip = step
        #     else:
        #         daily_mean_precip += step
        #     print(f"正在读取{nc_file},处理进度：{t/time_len:.2%}\n")
            
        # daily_mean_precip /= time_len


        # 将数据类型转换为16位整型，并设置Nodata值为-999
        daily_mean_precip = daily_mean_precip.fillna(-999).astype(np.int16)
        
        # 获取空间元数据信息
        transform = from_origin(ds.lon[0].values, ds.lat[0].values, abs(ds.lon[1] - ds.lon[0]), abs(ds.lat[1] - ds.lat[0]))
        width, height = daily_mean_precip.shape
        print("正在将结果保存为TIFF\n")
        # 设置输出文件名
        output_file = os.path.join(output_folder, f"{os.path.splitext(nc_file)[0]}_daily_mean.tif")

        # 将结果保存为GeoTIFF
        with rasterio.open(
            output_file,
            'w',
            driver='GTiff',
            height=height,
            width=width,
            count=1,
            dtype=rasterio.int16,
            crs='EPSG:4326',
            transform=transform,
            nodata=-999
        ) as dst:
            # 按块处理和写入GeoTIFF
            for i in range(0, height, block_size):
                for j in range(0, width, block_size):
                    # 读取当前块
                    window = ((i, min(i + block_size, height)), (j, min(j + block_size, width)))
                    data_block = daily_mean_precip[window[0][0]:window[0][1], window[1][0]:window[1][1]].values
                    
                    # 写入GeoTIFF
                    dst.write(data_block, 1, window=rasterio.windows.Window(j, i, data_block.shape[1], data_block.shape[0]))

        # 释放不再需要的数据
        del daily_mean_precip, ds
        print(f"已保存日均降水量TIFF文件: {output_file}\n")

# 使用示例
input_folder = 'E:\\1km分辨率土壤湿度\\10cm'  # 替换为NetCDF文件所在文件夹路径
output_folder = 'E:\\1km分辨率土壤湿度\\10cm'  # 替换为保存TIFF文件的文件夹路径
calculate_daily_mean_precipitation(input_folder, output_folder)
